from factor_graph import FactorGraph
import lietorch
from lietorch import SE3
import numpy as np
import torch


class DroidFrontend:

  def __init__(self, net, video, args):
    self.video = video
    self.update_op = net.update
    self.graph = FactorGraph(
        video, net.update, max_factors=48, upsample=args.upsample
    )

    # local optimization window
    self.t0 = 0
    self.t1 = 0

    # frontent variables
    self.is_initialized = False
    self.count = 0

    self.max_age = 25
    self.iters1 = 8
    self.iters2 = 4

    self.warmup = args.warmup
    self.beta = args.beta
    self.frontend_nms = args.frontend_nms
    self.keyframe_thresh = args.keyframe_thresh
    self.frontend_window = args.frontend_window
    self.frontend_thresh = args.frontend_thresh
    self.frontend_radius = args.frontend_radius

  def __update(self):
    """add edges, perform update"""
    self.count += 1
    self.t1 += 1

    if self.graph.corr is not None:
      self.graph.rm_factors(self.graph.age > self.max_age, store=True)

    self.graph.add_proximity_factors(
        self.t1 - 5,
        max(self.t1 - self.frontend_window, 0),
        rad=self.frontend_radius,
        nms=self.frontend_nms,
        thresh=self.frontend_thresh,
        beta=self.beta,
        remove=True,
    )

    self.video.disps[self.t1 - 1] = torch.where(
        self.video.disps_sens[self.t1 - 1] > 0,
        self.video.disps_sens[self.t1 - 1],
        self.video.disps[self.t1 - 1],
    )

    for itr in range(self.iters1):
      self.graph.update(
          None, None, use_inactive=True, viz_itr=None, use_mono=True
      )

    # set initial pose for next frame
    poses = SE3(self.video.poses)
    d = self.video.distance(
        [self.t1 - 3], [self.t1 - 2], beta=self.beta, bidirectional=True
    )

    if d.item() < self.keyframe_thresh:
      self.graph.rm_keyframe(self.t1 - 2)

      with self.video.get_lock():
        self.video.counter.value -= 1
        self.t1 -= 1

    else:
      for itr in range(self.iters2):
        self.graph.update(None, None, use_inactive=True)

    # set pose for next itration
    self.video.poses[self.t1] = self.video.poses[self.t1 - 1]
    self.video.disps[self.t1] = self.video.disps[self.t1 - 1].mean()

    # update visualization
    self.video.dirty[self.graph.ii.min() : self.t1] = True

  def __initialize(self):
    """initialize the SLAM system"""

    self.t0 = 0
    self.t1 = self.video.counter.value

    self.graph.add_neighborhood_factors(self.t0, self.t1, r=3)

    for itr in range(10):
      self.graph.update(1, use_inactive=True, use_mono=True, motion_only=True)

    self.graph.add_proximity_factors(
        0, 0, rad=2, nms=2, thresh=self.frontend_thresh, remove=False
    )

    for itr in range(10):
      self.graph.update(
          1, use_inactive=True, use_mono=True, motion_only=True, viz_itr=None
      )

    for itr in range(10):
      self.graph.update(
          1, use_inactive=True, use_mono=True, motion_only=False, viz_itr=None
      )

    # print("error ", error)
    # breakpoint()

    # self.video.normalize()
    self.video.poses[self.t1] = self.video.poses[self.t1 - 1].clone()
    self.video.disps[self.t1] = self.video.disps[self.t1 - 4 : self.t1].mean()

    # initialization complete
    self.is_initialized = True
    self.last_pose = self.video.poses[self.t1 - 1].clone()
    self.last_disp = self.video.disps[self.t1 - 1].clone()
    self.last_time = self.video.tstamp[self.t1 - 1].clone()

    with self.video.get_lock():
      self.video.ready.value = 1
      self.video.dirty[: self.t1] = True

    self.graph.rm_factors(self.graph.ii < self.warmup - 4, store=True)

  def __call__(self, final_=False):
    """main update"""
    if not self.is_initialized and final_ == True:
      self.__initialize()

    # do initialization
    if not self.is_initialized and self.video.counter.value == self.warmup:
      self.__initialize()

    # do update
    elif self.is_initialized and self.t1 < self.video.counter.value:
      self.__update()
